#include <cstdint>
#include <cstdio>
#include <random>
#include <systemc>
#include <tlm>
#include <tlm_utils/simple_initiator_socket.h>
#include <tlm_utils/simple_target_socket.h>

using namespace sc_core;
using namespace sc_dt;
using namespace tlm_utils;
using namespace std;

const sc_time DELAY = sc_time(2, SC_NS);

SC_MODULE(Master) {
    // Random number generator
    std::mt19937 gen;
    std::uniform_int_distribution<uint8_t> addr_dist;
    std::uniform_int_distribution<uint8_t> rwop_dist; // 0: read, 1: write
    // Initiator socket
    simple_initiator_socket<Master> socket;
    // DMI
    bool dmi_ptr_valid = false;
    tlm::tlm_dmi dmi_data;

    SC_CTOR(Master) : socket("initiator_socket"),
                      gen(0),
                      addr_dist(0, 64 - 1),
                      rwop_dist(0, 1),
                      dmi_ptr_valid(false) {
        // Register callbacks
        SC_THREAD(process);
        socket.register_invalidate_direct_mem_ptr(this, &Master::invalidate_direct_mem_ptr);
    }

    void process() {
        sc_time delay = DELAY;
        while (true) {
            uint8_t addr = addr_dist(gen) << 2;
            uint8_t rwop = rwop_dist(gen);
            uint32_t data = 0;

            // Prepare transaction
            tlm::tlm_generic_payload* payload = new tlm::tlm_generic_payload();
            tlm::tlm_command cmd = (rwop == 0) ? tlm::TLM_READ_COMMAND : tlm::TLM_WRITE_COMMAND;
            if (cmd == tlm::TLM_WRITE_COMMAND) { data = 0xFF000000 | addr; }

            // Use DMI if available
            printf("%s:\t%s REQ addr=0x%02X at time %s delay %s\n",
                    this->name(), (cmd == tlm::TLM_READ_COMMAND) ? "R" : "W",
                    addr, sc_time_stamp().to_string().c_str(), delay.to_string().c_str());
            if (dmi_ptr_valid) {
                auto dmi_addr = dmi_data.get_dmi_ptr() + (addr - dmi_data.get_start_address());
                if (cmd == tlm::TLM_READ_COMMAND) {
                    assert(dmi_data.is_read_allowed());
                    memcpy(&data, dmi_addr, 4);
                    delay = dmi_data.get_read_latency();
                }
                else if (cmd == tlm::TLM_WRITE_COMMAND) {
                    assert(dmi_data.is_write_allowed());
                    memcpy(dmi_addr, &data, 4);
                    delay = dmi_data.get_write_latency();
                }

                // Print transaction details
                printf("%s:\t%s RSP addr=0x%02X, data=0x%08X at time %s delay %s [DMI]\n",
                    this->name(), (cmd == tlm::TLM_READ_COMMAND) ? "R" : "W",
                    addr, data, sc_time_stamp().to_string().c_str(), delay.to_string().c_str());
            }
            else {
                // Set up the payload
                payload->set_command(cmd);
                payload->set_address(addr);
                payload->set_data_ptr(reinterpret_cast<unsigned char*>(&data));
                payload->set_data_length(sizeof(data));
                payload->set_byte_enable_ptr(0);
                payload->set_dmi_allowed(false);
                payload->set_response_status(tlm::TLM_INCOMPLETE_RESPONSE);

                // Send the transaction
                socket->b_transport(*payload, delay);

                // Check response status
                if (payload->is_response_error()) {
                    SC_REPORT_ERROR("TLM-2", "Response error from b_transport");
                }

                // Check for DMI possibility
                if (payload->is_dmi_allowed()) {
                    dmi_data.init();
                    dmi_ptr_valid = socket->get_direct_mem_ptr(*payload, dmi_data);
                }

                // Print transaction details
                printf("%s:\t%s RSP addr=0x%02X, data=0x%08X at time %s delay %s\n",
                    this->name(), (cmd == tlm::TLM_READ_COMMAND) ? "R" : "W",
                    addr, data, sc_time_stamp().to_string().c_str(), delay.to_string().c_str());
            }

            // Wait for the delay
            wait(delay);
        }
    }

    // TLM-2.0 backward DMI method
    void invalidate_direct_mem_ptr(sc_dt::uint64 start_range, sc_dt::uint64 end_range) {
        printf("%s: DMI pointer invalidated for [0x%llX, 0x%llX] at time %s\n",
               this->name(), start_range, end_range, sc_time_stamp().to_string().c_str());
        dmi_ptr_valid = false;
    }
};

SC_MODULE(Slave) {
    // Target socket
    simple_target_socket<Slave> socket;
    // Memory
    uint8_t mem[256];
    // Random number generator
    std::mt19937 gen;
    std::uniform_int_distribution<uint8_t> dist; // 0-255

    SC_CTOR(Slave) : socket("target_socket"), gen(1), dist(0, 32) {
        // Register callbacks
        SC_THREAD(invalidation_process);
        socket.register_b_transport(this, &Slave::b_transport);
        socket.register_get_direct_mem_ptr(this, &Slave::get_direct_mem_ptr);
        // Initialize memory
        uint32_t* base = reinterpret_cast<uint32_t*>(&mem[0]);
        for (int i = 0; i < 64; ++i) {
            base[i] = 0xAA000000 | (i << 2);
        }
        for (int i = 0; i < 256; i += 4) {
            printf("MEM[%02X] = 0x%08X\n", i, *((uint32_t*)&mem[i]));
        }
    }

    void b_transport(tlm::tlm_generic_payload& payload, sc_time& delay) {
        // Extract command, address, and data pointer
        tlm::tlm_command cmd = payload.get_command();
        uint64_t addr = payload.get_address();
        unsigned char* data_ptr = payload.get_data_ptr();
        unsigned int data_length = payload.get_data_length();

        // Check address and data length
        if (addr >= 256 || data_length != 4) {
            payload.set_response_status(tlm::TLM_ADDRESS_ERROR_RESPONSE);
            return;
        }

        // Operate on memory
        if (cmd == tlm::TLM_READ_COMMAND) {
            memcpy(data_ptr, &mem[addr], data_length);
            printf("%s:\tR     addr=0x%02X, data=0x%08X at time %s\n",
                   this->name(), (unsigned int)addr, *((uint32_t*)data_ptr), sc_time_stamp().to_string().c_str());
        } 
        else if (cmd == tlm::TLM_WRITE_COMMAND) {
            memcpy(&mem[addr], data_ptr, data_length);
            printf("%s:\tW     addr=0x%02X, data=0x%08X at time %s\n",
                   this->name(), (unsigned int)addr, *((uint32_t*)data_ptr), sc_time_stamp().to_string().c_str());
        }

        // Set delay
        delay = DELAY * 4;

        // Set response status
        payload.set_response_status(tlm::TLM_OK_RESPONSE);

        // Set DMI hint
        payload.set_dmi_allowed(true);
    }

    // TLM-2.0 forward DMI method
    bool get_direct_mem_ptr(tlm::tlm_generic_payload& payload, tlm::tlm_dmi& dmi_data) {
        // Set DMI properties
        dmi_data.allow_read_write();
        dmi_data.set_dmi_ptr(reinterpret_cast<unsigned char*>(&mem[0]));
        dmi_data.set_start_address(0);
        dmi_data.set_end_address(255);
        dmi_data.set_read_latency(DELAY);
        dmi_data.set_write_latency(DELAY);
        return true;
    }
    
    // Periodic DMI invalidation
    void invalidation_process() {
        while (true) {
            wait(DELAY * dist(gen));
            socket->invalidate_direct_mem_ptr(0, 255);
        }
    }
};

int sc_main(int argc, char **argv)
{
    // Instantiate modules
    Master master("master");
    Slave slave("slave");

    // Bind sockets
    master.socket.bind(slave.socket);

    // Start simulation
    sc_start(1000, SC_NS);
    return 0;
}
